#!/usr/bin/python3
# coding=utf-8
#
# Copyright (C) 2023-2024. Huawei Technologies Co., Ltd. All rights reserved.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# ===============================================================================

import numpy as np

softmaxTilingdt = []
softmaxTiling = []

perElementByBlk = 8

def add_tiling(value, dt):
    softmaxTiling.append(value)
    softmaxTilingdt.append(dt)

def getExponet(dimR):
    exponent = 0
    tmpDimR = dimR
    while True:
        if (perElementByBlk ** (exponent + 1)) > tmpDimR:
            return exponent
        exponent += 1

def getScaleFactor(dimR, exponent):
    return dimR // (perElementByBlk ** exponent)

def Align(size, alignFactor):
    if size % alignFactor != 0:
        return (size // alignFactor + 1) * alignFactor
    return size

def gen_golden_data_simple():
    oriRSize = 4096+ 5 
    oriASize = 4

    rloopLimit = oriRSize // 4096
    aloopLimit = 1

    rSize = 0
    if rloopLimit > 0:
        rSize = 4096
    
    aSize = oriASize / aloopLimit
    tailRSize = oriRSize - rSize * rloopLimit
    
    exponent = 0
    scaleFactor = 0
    if rSize > 0:
        exponent = getExponet(rSize)
        scaleFactor = getScaleFactor(rSize, exponent)
    
    tailExponent = 0
    tailScaleFactor = 0
    if tailRSize > perElementByBlk:
        tailExponent = getExponet(tailRSize)
        tailScaleFactor = getScaleFactor(tailRSize, tailExponent)
    
    mainR = tailScaleFactor * (perElementByBlk ** tailExponent)
    tailR = tailRSize - mainR
    tailRAlign8 = 0
    if tailR > 0:
        tailRAlign8 = Align(tailR, 8)

    add_tiling(perElementByBlk, ('perElementByBlk', np.uint32))
    add_tiling(oriRSize, ('oriRSize', np.uint32))
    add_tiling(oriASize, ('oriASize', np.uint32))
    add_tiling(rloopLimit, ('rloopLimit', np.uint32))
    add_tiling(aloopLimit, ('aloopLimit', np.uint32))
    add_tiling(rSize, ('rSize', np.uint32))
    add_tiling(aSize, ('aSize', np.uint32))
    add_tiling(tailRSize, ('tailRSize', np.uint32))
    add_tiling(exponent, ('exponent', np.uint32))
    add_tiling(scaleFactor, ('scaleFactor', np.uint32))
    add_tiling(tailExponent, ('tailExponent', np.uint32))
    add_tiling(tailScaleFactor, ('tailScaleFactor', np.uint32))
    add_tiling(mainR, ('mainR', np.uint32))
    add_tiling(tailR, ('tailR', np.uint32))
    add_tiling(tailRAlign8, ('tailRAlign8', np.uint32))

    reduceTempTBufSize = aSize * tailRAlign8
    if mainR > 0:
        reduceTempTBufSize = aSize * mainR
    elif rSize > 0 and aSize * mainR < aSize * scaleFactor * (perElementByBlk ** (exponent -1)):
        reduceTempTBufSize = aSize * scaleFactor * (perElementByBlk ** (exponent -1))
    
    add_tiling(reduceTempTBufSize, ('reduceTempTBufSize', np.uint32))
   
    #表示的是在Ub里这个src的大小
    srcTBufSize = aSize * (mainR + tailRAlign8)
    if rSize > 0:
        srcTBufSize = aSize * rSize
    add_tiling(srcTBufSize, ('srcTBufSize', np.uint32))


    tiling = np.array([tuple(softmaxTiling)], dtype=softmaxTilingdt)
    print(tiling)
    src = np.random.uniform(1, 10, [oriASize, oriRSize]).astype(np.float32)
    # dst = src
    softmaxMax0 = np.max(src, axis=1, keepdims=False)
    softmaxMax = np.max(src, axis=1, keepdims=True)
    print(softmaxMax0)
    softmaxMax1 = np.repeat(softmaxMax, perElementByBlk, axis=1)
    softmaxMax = np.repeat(softmaxMax, oriRSize, axis=1)
    src0 = src - softmaxMax
    src0 = np.exp(src0)
    # src0 = src
    softmaxSum0 = np.sum(src0, axis=1, keepdims=True)
    print(softmaxSum0)
    softmaxSum = np.repeat(softmaxSum0, perElementByBlk, axis=1)

    softmaxSum1 = np.repeat(softmaxSum0, oriRSize, axis=1)

    dst = src0 / softmaxSum1
    
    tiling.tofile("./input/tiling.bin")
    src.tofile("./input/src.bin")
    softmaxMax1.tofile("./output/golden_softmaxMax.bin")
    softmaxSum.tofile("./output/golden_softmaxSum.bin")
    dst.tofile("./output/golden_dst.bin")


if __name__ == "__main__":
    gen_golden_data_simple()
